from flask import Flask, request, jsonify
from flask_cors import CORS
from file import Documentsplit
from openai import OpenAI
import os
import uuid

app = Flask(__name__)
CORS(app)
app.secret_key = 'jdwiqjdjjijioiooj'

api_key = 'sk-bphodwvjqrzayelwjmbfvgmgwlzjbhpexxmvazkdemrlzyst'
api_url = 'https://api.siliconflow.cn/v1'
model = 'Qwen/Qwen3-235B-A22B'

client = OpenAI(api_key=api_key, base_url=api_url)

# 创建全局变量存储处理后的对象
processed_documents = {}


@app.route('/upload', methods=['POST'])
def upload_file():
    file_objects = []

    # 处理多文件上传
    if 'files' in request.files:
        files = request.files.getlist('files')
        for file in files:
            if file.filename:
                # 生成唯一文件名防止冲突
                unique_filename = f"{uuid.uuid4()}_{file.filename}"
                file_path = os.path.join('./uploads', unique_filename)

                # 确保上传目录存在
                os.makedirs(os.path.dirname(file_path), exist_ok=True)
                file.save(file_path)

                try:
                    # 处理文件并存储
                    ds = Documentsplit(file_path).load_file()
                    file_key = unique_filename
                    processed_documents[file_key] = ds
                    file_objects.append({
                        "file_key": file_key,
                        "original_filename": file.filename
                    })
                except Exception as e:
                    return jsonify({
                        "status": "error",
                        "message": f"Error processing file {file.filename}: {str(e)}"
                    }), 500

    # 生成会话ID标识一组文件
    session_id = str(uuid.uuid4())
    processed_documents[session_id] = {
        "files": file_objects,
        "timestamp": str(uuid.uuid4())  # 用于会话过期处理
    }

    return jsonify({
        "status": "success",
        "message": f"Processed {len(file_objects)} files",
        "session_id": session_id,
        "files": file_objects
    })


@app.route('/generate', methods=['POST'])
def generate_content():
    try:
        data = request.json
        session_id = data.get('session_id')  # 获取会话ID
        user_input = data.get('input', '')

        if not user_input:
            return jsonify({"status": "error", "message": "Input content is required"}), 400

        context = ""
        # 检查会话是否存在
        if session_id and session_id in processed_documents:
            session_data = processed_documents[session_id]
            files = session_data.get('files', [])

            # 从所有文件中提取上下文
            for file_obj in files:
                file_key = file_obj.get('file_key')
                if file_key and file_key in processed_documents:
                    ds = processed_documents[file_key]
                    docs = ds.similarity_search(user_input)
                    context += "\n".join([doc.page_content for doc in docs])

        # 处理无文件或无匹配上下文的情况
        if not context:
            context = "没有找到相关上下文信息"

        messages = [
            {'role': 'system', 'content': '你是一个基于知识库的AI助手，请根据以下上下文回答问题。可以简写，争取5秒内答完'},
            {'role': 'user', 'content': f"上下文:\n{context}\n\n问题:{user_input}"}
        ]

        response = client.chat.completions.create(
            model=model,
            messages=messages,
            stream=True
        )

        generated_text = ""
        for chunk in response:
            if chunk.choices and chunk.choices[0].delta.content:
                generated_text += chunk.choices[0].delta.content

        return jsonify({
            "status": "success",
            "input": user_input,
            "generated_content": generated_text,
            "context_length": len(context)  # 可选：返回上下文长度用于调试
        })

    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500


if __name__ == '__main__':
    app.run(port=5000)